import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import torch
import torch.optim as optim
from functools import partial

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei'] 
plt.rcParams['axes.unicode_minus'] = False

def complete_optimization_pipeline():
    """完整的优化流程示例：波士顿房价预测"""
    
    # 生成模拟数据
    np.random.seed(42)
    n_samples = 1000
    n_features = 5
    
    # 生成特征数据
    X = np.random.randn(n_samples, n_features)
    # 生成真实权重和偏置
    true_weights = np.array([2.5, -1.3, 0.8, 3.2, -0.5])
    true_bias = 1.7
    # 生成目标值（加入噪声）
    y = np.dot(X, true_weights) + true_bias + 0.1 * np.random.randn(n_samples)
    
    # 数据预处理
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
    
    # 不同优化算法比较
    optimizers = {
        'SGD': optim.SGD,
        'Momentum': partial(optim.SGD, momentum=0.9), # 使用partial固定momentum参数
        'Adam': optim.Adam,
        'RMSprop': optim.RMSprop
    }
    
    results = {}
    
    for name, optimizer_class in optimizers.items():
        print(f"\n训练 {name} 优化器...")
        
        # 模型定义
        model = torch.nn.Sequential(
            torch.nn.Linear(n_features, 1)
        )
        
        # 定义优化器
        optimizer = optimizer_class(model.parameters(), lr=0.01)
        criterion = torch.nn.MSELoss()
        
        # 训练模型
        train_losses = []
        test_losses = []
        
        for epoch in range(100):
            # 训练阶段
            model.train()
            optimizer.zero_grad()
            
            # 前向传播
            outputs = model(torch.FloatTensor(X_train))
            loss = criterion(outputs, torch.FloatTensor(y_train).unsqueeze(1))
            
            # 反向传播
            loss.backward()
            optimizer.step()
            
            # 评估阶段
            model.eval()
            with torch.no_grad():
                test_outputs = model(torch.FloatTensor(X_test))
                test_loss = criterion(test_outputs, torch.FloatTensor(y_test).unsqueeze(1))
            
            train_losses.append(loss.item())
            test_losses.append(test_loss.item())
        
        results[name] = {
            'train_losses': train_losses,
            'test_losses': test_losses,
            'final_params': list(model.parameters())
        }
        
        print(f"{name} 最终训练损失: {train_losses[-1]:.4f}")
    
    # 可视化比较结果
    plt.figure(figsize=(12, 8))
    
    for name, result in results.items():
        plt.plot(result['train_losses'], label=f'{name} (训练)')
        plt.plot(result['test_losses'], '--', label=f'{name} (测试)', alpha=0.7)
    
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('不同优化算法性能比较')
    plt.legend()
    plt.grid(True)
    plt.yscale('log')
    plt.show()
    
    return results

# 运行完整示例
results = complete_optimization_pipeline()